import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models

# 定义AlexNet模型
class AlexNet(nn.Module):
    def __init__(self, num_classes=2):  # 默认2类pizza和非pizza
        super(AlexNet, self).__init__()
        self.model = models.alexnet(pretrained=True)
        # 修改最后一层全连接层，根据实际类别数调整
        self.model.classifier[6] = nn.Linear(4096, num_classes)

    def forward(self, x):
        return self.model(x)

# 保存模型的函数
def save_model(model, path):
    torch.save(model.state_dict(), path)

# 加载模型的函数
def load_model(model, path):
    model.load_state_dict(torch.load(path))
    return model
